/**@file
   IO accessing functions.
   
   @par license 
   Copyright 2012 Lu, Ken. All rights reserved.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 
   @author: Lu, Ken (bluewish.ken.lu@gmail.com)
**/

#include <syspy/types.h>

#include "io.h"
#include "log.h"

#define KS_TAG             'KSTB'

/**
 * MMIO space read
 * 
 * @author klu2 (2012/4/20)
 * 
 * @param base              base address for MMIO range
 * @param length            length for MMIO range
 * @param output            output buffer
 * @param output_length     the length of output buffer
 * 
 * @retval STATUS_INVALID_PARAMETER 
 * @retval STATUS_INSUFFICIENT_RESOURCES 
 */
NTSTATUS
k_mmio_read (
    un_t  base,
    u32_t length,
    u8_t  *output,
    u32_t *output_length
    )
{
    NTSTATUS         status     = STATUS_SUCCESS;
    PHYSICAL_ADDRESS phyaddr    = {0};
    PMDL             pmdl       = NULL;
    u8_t*            mapped_address = NULL;
    PVOID            pbuffer    = NULL;
    u64_t            temp;

    k_func_enter();

    k_assert((output != NULL));
    k_assert(output_length != NULL);
    
    //
    // Do not support > 4G address
    //
    temp = (u64_t)base;
    if (temp >= 0x100000000ULL || length == 0 || *output_length == 0) {
        return STATUS_INVALID_PARAMETER;
    }

    //
    // Make sure the range is limited to 4G.
    //
    temp = (u64_t)base + (u64_t)length;
    if (temp > 0x100000000ULL) {
        length = (u32_t)(0x100000000ULL - base);
        k_warn("base + length > 4G, shorten length to 0x%x", length);
    }

    if (*output_length < length) {
        k_debug("output buffer too small, shorten length to 0x%x", *output_length);
        length = *output_length;
    }

    phyaddr.HighPart = 0;
    phyaddr.LowPart  = base & 0xFFFFFFFF;
    __try {
        __try {
            mapped_address = MmMapIoSpace(phyaddr, length, MmNonCached);
            if (mapped_address == NULL) {
                k_debug("Fail to map IO space!");
                status = STATUS_INSUFFICIENT_RESOURCES;
                __leave;
            }

            pmdl = IoAllocateMdl(mapped_address, length, FALSE, FALSE, NULL);
            if (pmdl == NULL) {
                k_debug("Fail allocate mdl for mapped IO address 0x%x!", mapped_address);
                status = STATUS_INSUFFICIENT_RESOURCES;
                __leave;
            }
            MmBuildMdlForNonPagedPool(pmdl);
            pbuffer = MmMapLockedPagesSpecifyCache(pmdl, KernelMode, MmNonCached, NULL, FALSE, NormalPagePriority);
            if (pbuffer == NULL) {
                k_debug("Fail locked page for mapped IO address 0x%x!", mapped_address);
                status = STATUS_INSUFFICIENT_RESOURCES;
                __leave;
            }

            memcpy(output, pbuffer, length);
            *output_length = length;

        } __except (EXCEPTION_EXECUTE_HANDLER) {
            k_debug("Exception when KSMmIoRead");
            *output_length = 0;
            status = STATUS_INSUFFICIENT_RESOURCES;
        }
    } __finally {
        if (pmdl != NULL) {
            IoFreeMdl(pmdl);
        }

        if (mapped_address != NULL) {
            MmUnmapIoSpace(mapped_address, length);
        }

    }
    
    k_func_leave();
    return status;
}

/**
 * Scan given signature buffer in range.
 * 
 * @author klu2 (2012/4/20)
 * 
 * @param start 
 * @param length 
 * @param scan_buffer 
 * @param scan_buffer_len 
 * @param first 
 * 
 * @return NTSTATUS 
 */
NTSTATUS
k_mmio_scan(
    un_t    start,                 ///< start address for scan range
    u32_t   length,                ///< length for scan range
    u8_t    *scan_buffer,          ///< the buffer for search
    u32_t   scan_buffer_len,       ///< the length of search buffer
    u32_t   *first
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    u32_t current, remaining, index;
    u32_t actualsize, workbuf_len, read_len;
    u8_t  *workbuf = NULL;
    u64_t temp;

    k_func_enter();

    k_assert(scan_buffer != NULL);
    k_assert(first != NULL);

    if ((length == 0) || (scan_buffer_len == 0) || 
        (scan_buffer == NULL) || (first == NULL)) {
        return STATUS_INVALID_PARAMETER;
    }
    
    //
    // Only support the scan buffer with length less than 0x1000
    //
    if (scan_buffer_len > 0x1000) {
        return STATUS_INVALID_PARAMETER;
    }

    temp = (u64_t)start;
    if (temp > 0x100000000ULL) {
        return STATUS_INVALID_PARAMETER;
    }

    temp = (u64_t)start + (u64_t)length;
    if (temp > 0x100000000ULL) {
        length = (u32_t)(0x100000000ULL - start);
        k_warn("Scan range exceed 4G, short length to 0x%x", length);
    }
    
    workbuf_len = length < 0x100000 ? length : 0x100000;
    workbuf = ExAllocatePoolWithTag(PagedPool, workbuf_len, KS_TAG);
    if (workbuf == NULL) {
        k_debug("Fail to allocate work buffer in 1M size, so try 64K.");
        workbuf_len = 0x10000;
        workbuf = ExAllocatePoolWithTag(PagedPool, workbuf_len, KS_TAG);
        if (workbuf == NULL) {
            status = STATUS_INSUFFICIENT_RESOURCES;
            goto exit;
        }
    }

    *first = 0;
    current = start;
    do {
        memset(workbuf, 0, workbuf_len);

        remaining = start + length - current;
        read_len  = remaining <= workbuf_len ? remaining : workbuf_len;
        
        actualsize = workbuf_len;
        status = k_mmio_read(current, read_len, workbuf, &actualsize);
        if (!NT_SUCCESS(status)) {
            k_debug("Fail to map MMIO range(0x%x, 0x%x)", current, read_len);
            break;
        }

        for (index = current; index + scan_buffer_len - current < actualsize; index ++) {
            if (memcmp(&workbuf[index - current], scan_buffer, scan_buffer_len) == 0) {
                *first = index;
                goto exit;
            }
        }

        current += actualsize;
    } while (current - start < length);

exit:
    if (workbuf != NULL) {
        ExFreePoolWithTag(workbuf, KS_TAG);
    }


    if (*first == 0) {
        status = STATUS_NOT_FOUND;
    }


    k_func_leave();
    return status;
}
